import fitz  # PyMuPDF
from typing import List, Dict, Tuple, Optional, Any
import re

class ImprovedPDFTableMerger:
    def __init__(self, pdf_path: str):
        self.doc = fitz.open(pdf_path)
        self.pages = [self.doc.load_page(i) for i in range(len(self.doc))]
        print(f"已加载PDF文件，共 {len(self.pages)} 页")

    def extract_page_elements(self, page_num: int) -> Dict[str, Any]:
        """
        提取页面中的所有元素：文本块、线条和区域信息
        """
        page = self.pages[page_num]
        elements = {
            "text_blocks": [],
            "lines": [],
            "tables": [],
            "page_rect": page.rect
        }

        # 提取文本块
        text_dict = page.get_text("dict")
        if "blocks" in text_dict:
            text_blocks = text_dict["blocks"]
            for block in text_blocks:
                if block["type"] == 0:  # 文本块
                    # 提取文本内容
                    text_lines = []
                    for line in block.get("lines", []):
                        line_text = " ".join([span["text"] for span in line.get("spans", [])])
                        text_lines.append(line_text)
                    
                    elements["text_blocks"].append({
                        "bbox": block["bbox"],
                        "text": "\n".join(text_lines),
                        "full_block": block
                    })

        # 提取线条和表格框线
        drawings = page.get_drawings()
        for drawing in drawings:
            if "items" in drawing:
                for item in drawing["items"]:
                    if item[0] == "l":  # 线条
                        # 线条坐标处理
                        if len(item) >= 2:
                            line_points = item[1]
                            if isinstance(line_points, (list, tuple)) and len(line_points) == 4:
                                x0, y0, x1, y1 = line_points
                                elements["lines"].append({
                                    "type": "line",
                                    "points": [(x0, y0), (x1, y1)],
                                    "bbox": (min(x0, x1), min(y0, y1), max(x0, x1), max(y0, y1))
                                })
                    elif item[0] == "re":  # 矩形
                        if len(item) >= 2:
                            rect_params = item[1]
                            if isinstance(rect_params, (list, tuple)) and len(rect_params) == 4:
                                x, y, w, h = rect_params
                                elements["lines"].append({
                                    "type": "rect",
                                    "bbox": (x, y, x + w, y + h)
                                })

        # 使用表格识别算法检测表格
        elements["tables"] = self._detect_tables_from_lines_and_text(elements["lines"], elements["text_blocks"])
        
        return elements

    def _detect_tables_from_lines_and_text(self, lines: List[Dict], text_blocks: List[Dict]) -> List[Dict]:
        """
        从线条和文本块中检测表格
        """
        tables = []
        
        # 提取水平线和垂直线
        horizontal_lines = []
        vertical_lines = []
        
        for line in lines:
            bbox = line["bbox"]
            # 水平线判断：高度很小
            if abs(bbox[3] - bbox[1]) < 2:
                horizontal_lines.append({
                    "y": (bbox[1] + bbox[3]) / 2,
                    "x0": bbox[0],
                    "x1": bbox[2],
                    "bbox": bbox
                })
            # 垂直线判断：宽度很小
            elif abs(bbox[2] - bbox[0]) < 2:
                vertical_lines.append({
                    "x": (bbox[0] + bbox[2]) / 2,
                    "y0": bbox[1],
                    "y1": bbox[3],
                    "bbox": bbox
                })
        
        # 按Y坐标排序水平线
        horizontal_lines.sort(key=lambda x: x["y"])
        # 按X坐标排序垂直线
        vertical_lines.sort(key=lambda x: x["x"])
        
        # 如果有足够的线条，可能构成表格
        if len(horizontal_lines) >= 2 and len(vertical_lines) >= 2:
            # 查找表格边界
            table_bbox = (
                min(vl["x"] for vl in vertical_lines),
                min(hl["y"] for hl in horizontal_lines),
                max(vl["x"] for vl in vertical_lines),
                max(hl["y"] for hl in horizontal_lines)
            )
            
            # 添加表格信息
            tables.append({
                "bbox": table_bbox,
                "horizontal_lines": horizontal_lines,
                "vertical_lines": vertical_lines,
                "confidence": 0.7  # 基于线条数量的置信度
            })
        
        return tables

    def detect_cross_page_tables(self, margin_threshold: float = 20.0) -> List[Tuple[int, int]]:
        """
        检测跨页表格，使用多种方法提高检测率
        """
        cross_page_tables = []
        
        for i in range(len(self.pages) - 1):
            page1_elements = self.extract_page_elements(i)
            page2_elements = self.extract_page_elements(i + 1)
            
            # 获取页面边界
            page1_rect = page1_elements["page_rect"]
            page2_rect = page2_elements["page_rect"]
            
            # 方法1: 检查页面1底部和页面2顶部的表格对齐
            page1_bottom_tables = []
            for table in page1_elements["tables"]:
                table_bbox = table["bbox"]
                # 检查表格是否靠近页面底部
                if page1_rect.y1 - table_bbox[3] <= margin_threshold:
                    page1_bottom_tables.append(table)
            
            page2_top_tables = []
            for table in page2_elements["tables"]:
                table_bbox = table["bbox"]
                # 检查表格是否靠近页面顶部
                if table_bbox[1] - page2_rect.y0 <= margin_threshold:
                    page2_top_tables.append(table)
            
            # 检查表格是否对齐
            for table1 in page1_bottom_tables:
                for table2 in page2_top_tables:
                    if self._are_rects_aligned(table1["bbox"], table2["bbox"], "horizontal", 10):
                        cross_page_tables.append((i, i + 1))
                        print(f"检测到跨页表格: 第 {i+1} 页和第 {i+2} 页")
                        print(f"  - 基于表格边界对齐")
                        break
                if (i, i + 1) in cross_page_tables:
                    break
            
            # 如果已经检测到表格，跳过其他检测方法
            if (i, i + 1) in cross_page_tables:
                continue
            
            # 方法2: 检查文本对齐和连续性
            text_continuity_score = self._check_text_continuity(
                page1_elements["text_blocks"],
                page2_elements["text_blocks"],
                page1_rect,
                page2_rect,
                margin_threshold
            )
            
            if text_continuity_score > 0.5:
                cross_page_tables.append((i, i + 1))
                print(f"检测到跨页表格: 第 {i+1} 页和第 {i+2} 页")
                print(f"  - 基于文本连续性，得分: {text_continuity_score:.2f}")
        
        return cross_page_tables

    def _are_rects_aligned(self, rect1, rect2, direction: str, tolerance: float) -> bool:
        """检查两个矩形是否在指定方向上对齐"""
        if direction == "horizontal":
            return (abs(rect1[0] - rect2[0]) < tolerance and
                    abs(rect1[2] - rect2[2]) < tolerance)
        elif direction == "vertical":
            return (abs(rect1[1] - rect2[1]) < tolerance and
                    abs(rect1[3] - rect2[3]) < tolerance)
        return False

    def _check_text_continuity(self, text_blocks1, text_blocks2, page1_rect, page2_rect, margin_threshold: float) -> float:
        """
        检查两页文本之间的连续性，改进版
        """
        if not text_blocks1 or not text_blocks2:
            return 0.0
        
        # 获取第一页底部和第二页顶部的文本
        bottom_texts = []
        for block in text_blocks1:
            if block["bbox"][3] > page1_rect.y1 - margin_threshold * 2:  # 底部区域
                bottom_texts.append(block)
        
        top_texts = []
        for block in text_blocks2:
            if block["bbox"][1] < page2_rect.y0 + margin_threshold * 2:  # 顶部区域
                top_texts.append(block)
        
        if not bottom_texts or not top_texts:
            return 0.0
        
        # 提取文本内容
        bottom_text = " ".join([block["text"] for block in bottom_texts])
        top_text = " ".join([block["text"] for block in top_texts])
        
        # 检查1: 数字序列连续性
        bottom_numbers = re.findall(r'\d+', bottom_text)
        top_numbers = re.findall(r'\d+', top_text)
        
        if bottom_numbers and top_numbers:
            try:
                # 检查最后一个底部数字和第一个顶部数字是否连续
                last_bottom = int(bottom_numbers[-1])
                first_top = int(top_numbers[0])
                if first_top == last_bottom + 1:
                    return 1.0  # 强连续信号
            except:
                pass
        
        # 检查2: 文本结构相似性
        bottom_lines = bottom_text.lower().split('\n')
        top_lines = top_text.lower().split('\n')
        
        # 计算每行的文本长度分布相似度
        bottom_lengths = [len(line.strip()) for line in bottom_lines if line.strip()]
        top_lengths = [len(line.strip()) for line in top_lines if line.strip()]
        
        # 如果有相似的行长度分布，可能是同一表格
        if len(bottom_lengths) > 0 and len(top_lengths) > 0:
            # 计算平均长度差异
            avg_bottom = sum(bottom_lengths) / len(bottom_lengths)
            avg_top = sum(top_lengths) / len(top_lengths)
            length_diff = abs(avg_bottom - avg_top) / max(avg_bottom, avg_top, 1)
            
            if length_diff < 0.3:  # 如果平均长度差异小于30%
                return 0.7
        
        # 检查3: 共同词汇
        bottom_words = set(re.findall(r'\w+', bottom_text.lower()))
        top_words = set(re.findall(r'\w+', top_text.lower()))
        common_words = bottom_words & top_words
        
        if common_words:
            # 计算共同词汇比例
            common_ratio = len(common_words) / max(len(bottom_words), len(top_words), 1)
            return min(0.6, common_ratio * 2)  # 最多返回0.6分
        
        return 0.0

    def extract_table_data(self, page_num: int, region: Optional[Tuple[float, float, float, float]] = None) -> List[List[str]]:
        """
        提取指定区域内的表格数据
        改进了文本提取和行列组织
        """
        elements = self.extract_page_elements(page_num)
        text_blocks = elements["text_blocks"]
        
        if region:
            # 只提取指定区域内的文本块
            x0, y0, x1, y1 = region
            text_blocks = [block for block in text_blocks
                          if (block["bbox"][0] <= x1 and  # 块的左边界在区域右边界左侧
                              block["bbox"][2] >= x0 and  # 块的右边界在区域左边界右侧
                              block["bbox"][1] <= y1 and  # 块的上边界在区域下边界上方
                              block["bbox"][3] >= y0)]    # 块的下边界在区域上边界下方
        
        if not text_blocks:
            return []
        
        # 按行组织文本块
        # 1. 按Y坐标排序文本块
        text_blocks.sort(key=lambda x: x["bbox"][1])
        
        # 2. 合并接近的行
        rows = []
        current_row = []
        last_y = None
        
        for block in text_blocks:
            bbox = block["bbox"]
            y_center = (bbox[1] + bbox[3]) / 2
            
            # 确定是否为新行的阈值，根据页面上文本行间距动态调整
            if last_y is None:
                current_row.append(block)
                last_y = y_center
            else:
                # 计算与上一行的垂直距离
                vertical_distance = abs(y_center - last_y)
                
                # 如果距离小于平均行高的一半，认为是同一行
                if vertical_distance < 10:  # 适当的阈值，可根据实际PDF调整
                    current_row.append(block)
                else:
                    # 完成当前行
                    if current_row:
                        # 按X坐标排序当前行的文本块
                        current_row.sort(key=lambda x: x["bbox"][0])
                        rows.append([b["text"] for b in current_row])
                    # 开始新行
                    current_row = [block]
                    last_y = y_center
        
        # 添加最后一行
        if current_row:
            current_row.sort(key=lambda x: x["bbox"][0])
            rows.append([b["text"] for b in current_row])
        
        return rows

    def merge_cross_page_tables(self, margin_threshold: float = 20.0):
        """
        合并检测到的跨页表格
        """
        cross_page_pairs = self.detect_cross_page_tables(margin_threshold)
        merged_tables = {}
        
        for page1, page2 in cross_page_pairs:
            print(f"\n处理跨页表格: 第 {page1+1} 页和第 {page2+1} 页")
            
            # 获取页面边界
            page1_rect = self.pages[page1].rect
            page2_rect = self.pages[page2].rect
            
            # 确定提取区域
            # 对于第一页，提取底部区域
            page1_region = (page1_rect.x0, page1_rect.y1 - margin_threshold * 3, 
                           page1_rect.x1, page1_rect.y1)
            
            # 对于第二页，提取顶部区域
            page2_region = (page2_rect.x0, page2_rect.y0, 
                           page2_rect.x1, page2_rect.y0 + margin_threshold * 3)
            
            # 提取表格数据
            table_part1 = self.extract_table_data(page1, page1_region)
            table_part2 = self.extract_table_data(page2, page2_region)
            
            # 去重处理：如果两部分有重叠的行，只保留一行
            if table_part1 and table_part2:
                # 检查最后几行和第一几行是否相似
                last_rows_page1 = table_part1[-2:] if len(table_part1) >= 2 else table_part1
                first_rows_page2 = table_part2[:2] if len(table_part2) >= 2 else table_part2
                
                # 简单的去重逻辑
                to_remove = 0
                for i, row1 in enumerate(reversed(last_rows_page1)):
                    for j, row2 in enumerate(first_rows_page2):
                        # 如果两行内容相似，则认为是重复
                        if self._rows_are_similar(row1, row2):
                            to_remove = len(last_rows_page1) - i
                            break
                    if to_remove > 0:
                        break
                
                if to_remove > 0:
                    print(f"  - 检测到 {to_remove} 行重叠内容，已去除重复")
                    table_part1 = table_part1[:-to_remove]
            
            # 合并表格
            merged_table = table_part1 + table_part2
            
            merged_tables[(page1, page2)] = merged_table
            print(f"合并后的表格有 {len(merged_table)} 行")
            
            # 显示前几行和后几行作为示例
            print("前几行数据:")
            for i, row in enumerate(merged_table[:3]):
                print(f"  行 {i+1}: {row}")
            
            if len(merged_table) > 6:
                print("...")
                print("后几行数据:")
                for i, row in enumerate(merged_table[-3:], len(merged_table)-2):
                    print(f"  行 {i}: {row}")
            elif len(merged_table) > 3:
                print("...")
                for i, row in enumerate(merged_table[3:], 4):
                    print(f"  行 {i}: {row}")
        
        return merged_tables
    
    def _rows_are_similar(self, row1: List[str], row2: List[str]) -> bool:
        """
        检查两行是否相似（可能是重复行）
        """
        # 如果长度差异太大，肯定不相似
        if abs(len(row1) - len(row2)) > 1:
            return False
        
        # 计算相似单元格的比例
        min_len = min(len(row1), len(row2))
        similar_cells = 0
        
        for i in range(min_len):
            # 简单的文本相似度检查
            text1 = row1[i].strip().lower()
            text2 = row2[i].strip().lower()
            
            # 如果文本完全相同，或者一个是另一个的子集
            if text1 == text2 or text1 in text2 or text2 in text1:
                similar_cells += 1
        
        # 如果80%以上的单元格相似，则认为行相似
        return similar_cells / min_len >= 0.8

    def save_merged_tables_to_csv(self, merged_tables, output_dir: str = "."):
        """
        将合并后的表格保存为CSV文件
        """
        import os
        import csv
        
        for (page1, page2), table_data in merged_tables.items():
            if table_data:
                csv_filename = os.path.join(output_dir, f"merged_table_pages_{page1+1}_{page2+1}.csv")
                
                with open(csv_filename, 'w', newline='', encoding='utf-8') as csvfile:
                    writer = csv.writer(csvfile)
                    for row in table_data:
                        writer.writerow(row)
                
                print(f"已将合并后的表格保存到: {csv_filename}")

# 测试脚本
if __name__ == "__main__":
    import sys
    import os
    
    if len(sys.argv) < 2:
        print("用法: python improved_table_detection.py <pdf文件路径> [输出目录]")
        print("示例: python improved_table_detection.py 'D:\\工作\\02-文档智能\\05-测试数据\\01-跨页表格\\跨页表格01.pdf'")
        sys.exit(1)
    
    pdf_path = sys.argv[1]
    output_dir = sys.argv[2] if len(sys.argv) > 2 else os.path.dirname(pdf_path)
    
    if not os.path.exists(pdf_path):
        print(f"错误: 找不到PDF文件 '{pdf_path}'")
        sys.exit(1)
    
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    try:
        print(f"开始处理PDF文件: {pdf_path}")
        merger = ImprovedPDFTableMerger(pdf_path)
        
        # 调整参数以提高检测率
        print("正在检测跨页表格...")
        merged_tables = merger.merge_cross_page_tables(margin_threshold=30.0)  # 增加阈值以捕获更多可能的跨页表格
        
        if merged_tables:
            print(f"\n成功检测并合并了 {len(merged_tables)} 个跨页表格")
            
            # 保存结果
            merger.save_merged_tables_to_csv(merged_tables, output_dir)
        else:
            print("\n未检测到跨页表格。尝试以下方法:\n" \
                 "1. 调整margin_threshold参数\n" \
                 "2. 检查PDF文件是否包含表格\n" \
                 "3. 确认表格是否确实跨页")
    
    except Exception as e:
        print(f"处理PDF时发生错误: {e}")
        import traceback
        traceback.print_exc()